""" Plot density of states (DOS) stored in a VASP-DOSCAR file.
Original version: GHB, 20170711.
Update 2020-04-29: The first item of each data set in DOSCAR is incorrect; 
    it is skipped.
"""
import numpy as np
import copy

def get_dos(doscar="DOSCAR",
    partial=True,
    include_accumulated_dos=False) :
    """Read density of states from a DOSCAR file.

    The format of DOSCAR files varies with spin-polarization, largest angular
    quantum number (lmax, typical values are 2 for d, and 3 for f), components
    of projection (l or lm), and site decomposition (i.e., whether
    decomposition is done on each atom).

    This function tries to smartly work out the correct format. See 
    get_dos_manual_setting() for manual setting of the format.
    The guessing process is as follows:
    1. The sixth line of DOSCAR contains the number of spin-components.
    2. If there are more nonempty lines following the total DOS, then 
       the densities of states are site- and orbital-projected; that is
       partial DOS are available.
    3. The number of fields in each line of partial DOS map uniquely to 
       the format as follows.

    spin | orbital | lmax | n_f | notes
    1    | 'l'     | 2(d) | 4   | energy, s, p, d
    1    | 'l'     | 3(f) | 5   | energy, s, p, d, f
    1    | 'lm'    | 2(d) | 10  | energy, s, px, py, pz, ..., d3 
    1    | 'lm'    | 3(f) | 17  | energy, s, px, py, pz, ..., f5
    2    | 'l'     | 2(d) | 7   | energy, s-up, s-down, p-up, ..., d-down
    2    | 'l'     | 3(f) | 9   | energy, s-up, s-down, p-up, ..., f-down
    2    | 'lm'    | 2(d) | 19  | energy, s-up, s-down, px-up, ..., d3-down
    2    | 'lm'    | 3(f) | 33  | energy, s-up, s-down, px-up, ..., f5-down

    It is thus possible to deduce the settings by the number of fields.
    CAVEATS: VASP may change the format of DOSCAR files.

    Arguments:
    doscar: (string) path of DOSCAR file; default 'DOSCAR', case sensitive.
    partial: (boolean) whether to return partial dos; defulat True, but if 
        there are no projected DOS in the file, then None is
        returned in place of partial_dos. If partial is set to False, then
        only total DOS is returned. 
    include_accumulated_dos: (boolean) if True, the accumulated DOS for each 
        spin component is returned as the result's last element. The 
        accumulated DOS is usually not used, so the default is False (i.e., 
        do not return accumulated DOS.)

    Returns:
    When partial==True and include_accumulated_dos==False:
        a tuple (energy, total_dos, partial_dos)
    When partial==True and include_accumulated_dos==True:
        a tuple (energy, total_dos, partial_dos, accumulated_dos)
    When partial==False and include_accumulated_dos==False:
        a tuple (energy, total_dos)
    When partial==False and include_accumulated_dos==True:
        a tuple (energy, total_dos, accumulated_dos)
    in which
    energy: (one-dimensional N-element array) energies.
        N is equal to NEDOS-1 in the INCAR file.
    total_dos: (M by N array) total density of states. 
        M is equal to the number of spin components.
        if spin==1, then total_dos is a one-dimensional N-element array.
        if spin==2, then total_dos contains two one-dimensional N-element
        arrays, one for spin-up, and the other for spin-down.
    partial_dos: (None, or A by B by N numpy.array) site- and orbital-projected 
        density of states. A is the number of atoms; B is variable depending 
        on spin, lmax, and component:
            if component==None, then partial_dos = None
            else if spin==1, lmax==2, and orbital=='l', then B=3
            else if spin==1, lmax==2, and orbital=='lm', then B=9
            else if spin==1, lmax==3, and orbital=='l', then B=4
            else if spin==1, lmax==3, and orbital=='lm', then B=16
            else if spin==2, B is doubled compared to those for spin==1.
    """
    # Reads in the DOSCAR file, and save all lines.
    doscar_fn = doscar
    f = open(doscar_fn, "r")
    all_lines = f.readlines()
    f.close()
    
    # Parses several important parameters
    n_atoms = int(all_lines[0].split()[0])   # First field of line 0

    e_max, e_min, n_dos, e_fermi, w = [float(x) for x in all_lines[5].split()]
    n_dos = int(n_dos)

    # Guesses the number of spin components, n_spin.
    n_fields = len(all_lines[6].split())
    if n_fields==3 : 
        n_spin = 1
    elif n_fields==5 :
        n_spin = 2
    else :
        print("Line 6 of DOSCAR should contain 3 or 5 fields.")
        print("The actual number of fields is", n_fields)
        print("I stop parsing this DOSCAR file.")
        raise ValueError("Unexpected number of fields of total DOS.")

    # Parses total DOS and accumulated DOS
    energy_array, dos_total, dos_accumulated = _total_dos(
            all_lines, n_dos, n_spin, e_fermi, include_accumulated_dos)

    # Returns total DOS (and accumulated DOS if specified) if partial==False.
    if not partial: 
        if not include_accumulated_dos:
            return (energy_array, dos_total)
        else :
            return (energy_array, dos_total, dos_accumulated)
    # otherwise, continue parsing the DOSCAR file.

    # Does the file contain partial dos?
    line_start = 8+n_dos
    if len(all_lines)>line_start :
        n_fields = len(all_lines[line_start].split())
        if n_fields==4 :
            component = 'l'
            lmax = 2
            assert n_spin==1
        elif n_fields==5 :
            component = 'l'
            lmax = 3 
            assert n_spin==1
        elif n_fields==7 :
            component = 'l'
            lmax = 2
            assert n_spin==2
        elif n_fields==9 :
            component = 'l'
            lmax = 3
            assert n_spin==2
        elif n_fields==10 :
            component = 'lm'
            lmax = 2
            assert n_spin==1
        elif n_fields==17 :
            component = 'lm'
            lmax = 3
            assert n_spin==1
        elif n_fields==19 :
            component = 'lm'
            lmax = 2
            assert n_spin==2
        elif n_fields==33 :
            component = 'lm'
            lmax = 3
            assert n_spin==2
        else :
            print("*Error: number of fields should be 4,5,7,9,10,17,19, or 33.")
            print("        but in DOSCAR it was", n_fields)
            raise ValueError("Wrong number of fields of partial DOS in DOSCAR.")

        dos_partial = _projected_dos(
            all_lines, n_atoms, n_dos, lmax, n_spin, component)
    else:
        dos_partial = []

    if include_accumulated_dos: 
        return (energy_array, dos_total, dos_partial, dos_accumulated)
    else :
        return (energy_array, dos_total, dos_partial)
    

def get_dos_manual_setting(doscar="DOSCAR", 
    n_spin=1, 
    lmax=2, 
    partial=True,
    component=None,
    include_accumulated_dos=False):
    """ Read density of states from a DOSCAR file.

    The format of DOSCAR files varies with spin-polarization, largest angular
    quantum number (l, typical values are 2 for d, and 3 for f), components
    of projection (l or lm), and site decomposition (i.e., whether 
    decomposition is done on each atom).

    Arguments:
    doscar: (string) path of DOSCAR file, default 'DOSCAR', case sensitive.
    n_spin: (integer, valid values are 1 and 2) number of spin components, 
        default 1.
    lmax: (integer, valid values are 2 and 3) largest quantum number of 
        angular momentum; s, p, d, and f are 0, 1, 2, and 3, respectively,
        default 2 (d).
    partial: (boolean) whether to return partial dos. If partial is set to 
        True but there are no projected DOS in the file, then None is
        returned in place of partial_dos. If partial is set to False, then
        only total DOS is returned. 
    component: (None or string, valid values are None, "", 'l' and 'lm')
        if set to 'l', decompositions are done on angular orbitals only;
        if set to 'lm', decompositions are done on magnetization orbitals
        also; if set to None or "", partial DOS are not parsed nor returned.
        In VASP's input file INCAR, LORBIT=10 corresponds to the 'l'
        setting, and LORBIT=11 corresponds to the 'lm' setting. 
        An s orbital has zero magnetic quantum number,
        A p orbital has three magnetic quantum numbers: px, py, and pz;
        A d orbital has five magnetic quantum numbers: dxy, dxz, dz2, 
            dyz, and dx2-y2; 
        An f orbital has seven magnetic quantum numbers: fx(x2-3y2),
            fy(x2-z2), fxz2, fz3, fyz2, fxyz, and fy(3x2-y2).
    include_accumulated_dos: (boolean) if True, the accumulated DOS for each 
        spin component. The accumulated DOS is usually not used, so the 
        default is False (i.e., do not return accumulated DOS).

    Returns:
    When partial==True and include_accumulated_dos==False:
        a tuple (energy, total_dos, partial_dos)
    When partial==True and include_accumulated_dos==True:
        a tuple (energy, total_dos, partial_dos, accumulated_dos)
    When partial==False and include_accumulated_dos==False:
        a tuple (energy, total_dos)
    When partial==False and include_accumulated_dos==True:
        a tuple (energy, total_dos, accumulated_dos)
    in which
    energy: (one-dimensional N-element array) energies.
        N is equal to NEDOS-1 in the INCAR file.
    total_dos: (M by N array) total density of states. 
        M is equal to the number of spin components.
        if n_spin==1, then total_dos is a one-dimensional N-element array.
        if n_spin==2, then total_dos contains two one-dimensional N-element
        arrays, one for spin-up, and the other for spin-down.
    partial_dos: (None, or A by B by N numpy.array) site- and orbital-projected 
        density of states. A is the number of atoms; B is variable depending 
        on n_spin, lmax, and component:
            if component==None, then partial_dos = None
            else if n_spin==1, lmax==2, and orbital=='l', then B=3
            else if n_spin==1, lmax==2, and orbital=='lm', then B=9
            else if n_spin==1, lmax==3, and orbital=='l', then B=4
            else if n_spin==1, lmax==3, and orbital=='lm', then B=16
            else if n_spin==2, B is doubled compared to those for n_spin==1.
    When component is invalid, partial DOS will not be parsed or returned.
    """
    # Processes arguments
    doscar_fn = doscar
    
    if n_spin not in (1, 2) :
        raise ValueError("Number of spin components should be 1 or 2.")

    if lmax not in (2, 3) :
        raise ValueError("Maximum angular orbital should be 2 or 3.")

    if component is None:
        decomposition = False
    else:
        if component not in ('l', 'lm'):
            raise ValueError("component should be None, 'l', or 'lm'.")
        else:
            decomposition = True

    # Read DOSCAR and save all lines in this file
    f = open(doscar_fn, "r")
    all_lines = f.readlines()
    f.close()
    
    # Read several important parameters
    n_atoms = int(all_lines[0].split()[0])   # First number of line 1
    e_max, e_min, n_dos, e_fermi, s = [float(x) for x in all_lines[5].split()]
    n_dos = int(n_dos)

    assert n_spin==int(s)

    # Total dos is always present.
    energy_array, dos_total, dos_accumulated = _total_dos(
            all_lines, n_dos, n_spin, e_fermi, include_accumulated_dos)
    # Partial dos may not exist or needed.
    if partial :
        if decomposition :
            dos_partial = _projected_dos(
                all_lines, n_atoms, n_dos, lmax, spin, component)
        else :
            dos_partial = []
    
    # return (energy, dos_total) if partial==False.
    # otherwise, continue parsing the DOSCAR file.
    if not partial: 
        if include_accumulated_dos:
            return (energy_array, dos_total, dos_accumulated)
        else :
            return (energy_array, dos_total)
    else:
        if include_accumulated_dos:
            return (energy_array, dos_total, dos_partial, dos_accumulated)
        else :
            return (energy_array, dos_total, dos_partial)


def _total_dos(lines_in_DOSCAR, N_DOS, N_Spin, E_Fermi, Include_Acc):
    """ Process total DOS which is always present in a DOSCAR."""
    energy_array = []
    if N_Spin==1:
        if Include_Acc :
            dos_total = []
            dos_accumulated = []
            for line in lines_in_DOSCAR[6+1 : 6+N_DOS] :
                energy, dos_curr, dos_acc = [float(x) for x in line.split()]
                energy_array.append(energy-E_Fermi)
                dos_total.append(dos_curr)
                dos_accumulated.append(dos_acc)
        else :
            dos_total = []
            for line in lines_in_DOSCAR[6+1 : 6+N_DOS] :
                energy, dos_curr, dos_acc = [float(x) for x in line.split()]
                energy_array.append(energy-E_Fermi)
                dos_total.append(dos_curr)
    elif N_Spin==2:
        if Include_Acc :
            dos_total = [[], []]
            dos_accumulated = [[], []]
            for line in lines_in_DOSCAR[6+1 : 6+N_DOS] :
                energy, dos_up_curr, dos_down_curr, dos_up_acc, dos_down_acc=[
                    float(x) for x in line.split()]
                energy_array.append(energy-E_Fermi)
                dos_total[0].append(dos_up_curr)
                dos_total[1].append(dos_down_curr)
                dos_accumulated[0].append(dos_up_acc)
                dos_accumulated[1].append(dos_down_acc)
        else :
            dos_total = [[], []]
            for line in lines_in_DOSCAR[6+1 : 6+N_DOS] :
                energy, dos_up_curr, dos_down_curr, dos_up_acc, dos_down_acc=[
                    float(x) for x in line.split()]
                energy_array.append(energy-E_Fermi)
                dos_total[0].append(dos_up_curr)
                dos_total[1].append(dos_down_curr)
    else:
        print("Error: The number of spin components should be 1 or 2.")
        print("       But in DOSCAR it was", n_spin, ".")
        raise ValueError("Number of spin components should be 1 or 2.")
    
    # Returns total DOS always, return accumulated DOS if asked so.
    if Include_Acc :
        return (energy_array, dos_total, dos_accumulated)
    else :
        return (energy_array, dos_total, None)


def _projected_dos(lines_in_DOSCAR, N_Atoms, N_DOS, L_Max, N_Spin, Component):
    """ Process site- and orbital-projected DOS)."""
    if N_Spin==1:
        if L_Max==2:   # highest orbital: d
            if Component=='l' :
                p = [[], [], [], []]   # s, p, d
            elif Component=='lm' :
                # s, p_x, p_y, p_z, d-2, d-1, d0, d1, d2
                p = [[],    # s
                     [],[],[],    # p
                     [],[],[],[],[],  ]    # d
            else:
                pass
        elif L_Max==3:   # highest orbital: f
            if Component=='l' :
                p = [[], [], [], []]   # s, p, d, f
            elif Component=='lm' :
                # 1 s, 3 p, 5 d, 7 f
                p = [[],  # s
                     [],[],[],  # p
                     [],[],[],[],[],   # d
                     [],[],[],[],[],[],[], ]   # f
            else:
                pass
        else: 
            pass
    elif N_Spin==2:
        if L_Max==2:   # highest orbital: d
            if Component=='l' :
                # s-up, s-down, p-up, p-down, d-up, d-down
                p = [[],[], [],[], [],[]]   # s, p, d
            elif Component=='lm' :
                # 2 s, 6 p, 10 d = 18
                p = [[],[],  # s-up, s-down 
                     [],[], [],[], [],[], # px-up, px-down, ...
                     [],[], [],[], [],[], [],[], [],[]  # d-2-up, d-2-down, ...
                ]
            else:
                pass
        elif L_Max==3:   # highest orbital: f
            if Component=='l':
                # s-up, s-down, p-up, p-down, d-up, d-down, f-up, f-down
                p = [[],[], [],[], [],[], [],[]]   
            elif Component=='lm':
                # 2 s, 6 p, 10 d, 14 f
                p = [[],[],  # s-up, s-down 
                     [],[], [],[], [],[], # px-up, px-down, ...
                     [],[], [],[], [],[], [],[], [],[], #d-2-up, d-2-down, ...
                     [],[], [],[], [],[], [],[], [],[], [],[], [],[], # f
                ]
            else:
                pass
        else:
            pass
    else:
        pass

    dos_partial = []
    for atom_index in range(N_Atoms):
        dos_partial.append(copy.deepcopy(p))

        line_start = 8+N_DOS + atom_index * (N_DOS+1)
        line_end   = line_start + N_DOS - 1
        for line in lines_in_DOSCAR[line_start : line_end] :
            line_buffer = [float(x) for x in line.split()[1:]]
            for i, d in enumerate(line_buffer):
                dos_partial[atom_index][i].append(d)

    return dos_partial

if __name__ =="__main__":
    import matplotlib.pyplot as plt
    F = plt.figure(figsize=(4,3))

    energy_array, dos_array = get_dos1()
    plt.plot(energy_array, dos_array, 'r-')
    plt.xlim(-20.0, 10.0)
    plt.ylim(0, 200)
    plt.xlabel("Energy (eV)")
    plt.ylabel(r"DOS (eV$^{-1}\,$Cell$^{-1}$)")
    plt.subplots_adjust(left=0.21, bottom=0.27, right=0.95)
    plt.savefig("dos1.png", dpi=300)
    plt.show()

